import warnings
import numpy as np
import matplotlib.pyplot as plt

from astropy.io import fits
from astropy.stats import sigma_clipped_stats

# photutils & astropy
from photutils.detection import DAOStarFinder
from photutils.psf import PSFPhotometry, IterativelySubtractedPSFPhotometry, IntegratedGaussianPRF, SourceGrouper
from photutils.background import LocalBackground, MMMBackground
from photutils.aperture import CircularAperture, aperture_photometry


def main():
    """
    主函数，演示如何使用 photutils 进行：
    1. 基础 PSF 拟合光度
    2. 迭代式 PSF 拟合光度
    3. 残差图像检查
    4. 与简易孔径光度对比
    """

    # -----------------------------
    # 1. 读取 FITS 文件数据
    # -----------------------------
    fits_filename = '/mnt/7b21f1e1-eb25-4cd5-bdb5-06d7d82fa253/Temp/force_photmetry/images/1278305/mb_sc_t1278305-19687_u_20250101132605_194_sciimg.fits'  # 替换为你的 FITS 文件路径
    with fits.open(fits_filename) as hdul:
        data = hdul[0].data.astype(float)

    # -----------------------------
    # 2. 基础图像统计量 & 星源检测
    # -----------------------------
    mean_val, med_val, std_val = sigma_clipped_stats(data, sigma=3.0)
    print(f"Image statistics: mean={mean_val:.3f}, median={med_val:.3f}, std={std_val:.3f}")

    # 使用 DAOStarFinder 进行源检测
    fwhm_guess = 3.0  # 假设星点 FWHM ~ 3 pix
    threshold = 5.0 * std_val
    dao_finder = DAOStarFinder(fwhm=fwhm_guess, threshold=threshold)

    # 在 data - median 上执行检测
    sources = dao_finder(data - med_val)
    if (sources is None) or (len(sources) == 0):
        raise ValueError("未检测到星源，请检查 FITS 数据或调整 fwhm/threshold 等参数！")
    print(f"检测到 {len(sources)} 个星源")

    # 添加 x_0, y_0 以便 PSFPhotometry 进行初始拟合
    sources['x_0'] = sources['xcentroid']
    sources['y_0'] = sources['ycentroid']

    # -----------------------------
    # 3. 定义 PSF 模型 & 拟合区域
    # -----------------------------
    # 使用 IntegratedGaussianPRF
    # 若你知道 FWHM ~ 3.0, 则 sigma ~ FWHM / 2.355 ~ 1.27. 这里初值取1.5也可以
    psf_model = IntegratedGaussianPRF(sigma=1.5)
    fit_shape = (11, 11)  # PSF 拟合时使用的局部裁剪区大小

    # -----------------------------
    # 4. 星源分组器（可选）
    # -----------------------------
    # 若担心相邻星源的 PSF 裁剪区会相互干扰，可通过 SourceGrouper 分组后同时拟合
    grouper = SourceGrouper(min_separation=2.0)

    # -----------------------------
    # 5. 局部背景估计
    # -----------------------------
    # 使用 MMMBackground 算法做局部背景估计，局部环半径 5-10 像素
    bkgstat = MMMBackground()
    localbkg_estimator = LocalBackground(
        inner_radius=5,
        outer_radius=10,
        bkg_estimator=bkgstat
    )

    # -----------------------------
    # 6. (A) 基础 PSF 拟合光度
    # -----------------------------
    psfphot = PSFPhotometry(
        psf_model=psf_model,
        fit_shape=fit_shape,
        finder=dao_finder,           # 可选，让 PSFPhotometry 内部也能再次检测源
        grouper=grouper,
        aperture_radius=4,           # 仅用于初步光度提取，与真实 FWHM 大小接近即可
        localbkg_estimator=localbkg_estimator
    )

    # 执行PSF光度测量
    with warnings.catch_warnings():
        # photutils 的一些运算可能触发 RuntimeWarning，可选择忽略
        warnings.simplefilter("ignore", RuntimeWarning)
        phot_table = psfphot(data)
    print("\n===== 基础PSF拟合光度结果 =====\n", phot_table)

    # -----------------------------
    # 6. (B) 计算残差图像
    # -----------------------------
    # 获取模型图像并手动计算残差
    model_image = psfphot.make_model_image(data, psf_shape=fit_shape, include_localbkg=True)
    residual_image = data - model_image

    # -----------------------------
    # 6. (C) 迭代式 PSF 拟合 (拥挤场景)
    # -----------------------------
    # 如果星场较拥挤，可以考虑 IterativelySubtractedPSFPhotometry
    iter_phot = IterativelySubtractedPSFPhotometry(
        finder=dao_finder,
        group_maker=grouper,
        psf_model=psf_model,
        fit_shape=fit_shape,
        aperture_radius=4,
        localbkg_estimator=localbkg_estimator,
        niters=3  # 迭代次数，可以根据需要增大
    )

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        iter_table = iter_phot(data)

    print("\n===== 迭代式PSF拟合光度结果 =====\n", iter_table)
    iter_residual = iter_phot.get_residual_image()

    # -----------------------------
    # 7. 结果可视化：残差图对比
    # -----------------------------
    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
    im0 = ax[0].imshow(data, origin='lower', cmap='gray', vmin=med_val-2*std_val, vmax=med_val+5*std_val)
    ax[0].set_title("原始图像")
    fig.colorbar(im0, ax=ax[0], fraction=0.046, pad=0.04)

    im1 = ax[1].imshow(residual_image, origin='lower', cmap='gray', vmin=-5*std_val, vmax=5*std_val)
    ax[1].set_title("单次PSF拟合残差")
    fig.colorbar(im1, ax=ax[1], fraction=0.046, pad=0.04)

    im2 = ax[2].imshow(iter_residual, origin='lower', cmap='gray', vmin=-5*std_val, vmax=5*std_val)
    ax[2].set_title("迭代式PSF拟合残差")
    fig.colorbar(im2, ax=ax[2], fraction=0.046, pad=0.04)

    plt.tight_layout()
    plt.show()

    # -----------------------------
    # 8. 与简易孔径光度对比 (Aperture Photometry)
    # -----------------------------
    # 在不太拥挤的场景，可以做孔径光度来大致对比
    # 这里直接使用初步检测到的 sources
    positions = np.transpose((sources['xcentroid'], sources['ycentroid']))
    aperture_radius = 4
    aperture = CircularAperture(positions, r=aperture_radius)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        aperture_phot = aperture_photometry(data, aperture)

    print("\n===== 孔径光度测量结果（示例）=====\n", aperture_phot)

    # 可以做简单的散点对比： aperture_sum vs. PSF 拟合得到的 flux_0
    # 这里以单次PSF结果为例
    # 注意：需要保证按相同星的顺序匹配，这里只演示“近似”对照
    # 更严格做法是再次 cross-match x_0, y_0
    plt.figure(figsize=(6, 5))
    plt.scatter(aperture_phot['aperture_sum'], phot_table['flux_0'], alpha=0.5)
    plt.xlabel("Aperture Flux (r=4)")
    plt.ylabel("PSF Fitted Flux")
    plt.title("比较：孔径光度 vs. PSF拟合光度")
    plt.grid(True)
    plt.show()

    # -----------------------------
    # 9. 后续：光度定标 & 科学分析
    # -----------------------------
    print("\n下一步：根据已知的零点 (ZP)、大气消光、颜色项等，进行星等转化：")
    print("  m = -2.5 * log10(flux) + ZP (若flux单位与ZP定义一致)")
    print("更多细节可参见 photutils 文档或专业文献。")


if __name__ == "__main__":
    main()
